import numpy as np

def calculate_iou(pred, target, num_classes, ignore_index=255):
    """
    计算每个类别的 IoU，并返回 mIoU (mean IoU)。
    
    参数：
        pred (torch.Tensor): 预测的类别标签，形状为 (N, H, W)。
        target (torch.Tensor): 真实的类别标签，形状为 (N, H, W)。
        num_classes (int): 类别总数。
        ignore_index (int): 忽略的类别标签，比如 255 表示不关注区域。
        
    返回：
        float: 平均 IoU (mIoU)。
    """
    # 先展平
    pred = pred.view(-1)
    target = target.view(-1)

    # 生成 mask，过滤掉 ignore_index 位置
    valid_mask = (target != ignore_index)
    pred = pred[valid_mask]
    target = target[valid_mask]

    iou_list = []
    for cls in range(num_classes):
        # 计算当前 cls 的交集和并集
        pred_inds = (pred == cls)
        target_inds = (target == cls)
        intersection = (pred_inds & target_inds).sum().float().item()
        union = (pred_inds | target_inds).sum().float().item()

        if union == 0:
            # 如果某个类别在这一批里不存在，可以忽略也可以视作0，通常忽略用 nan
            iou = float('nan')
        else:
            iou = intersection / union

        iou_list.append(iou)

    return np.nanmean(iou_list)

def mean_pixel_accuracy(pred, target, num_classes, ignore_index=255):
    """
    计算每个类别的像素准确率，并返回 mPA (mean Pixel Accuracy)。
    
    参数：
        pred (torch.Tensor): 预测的类别标签，形状为 (N, H, W)。
        target (torch.Tensor): 真实的类别标签，形状为 (N, H, W)。
        num_classes (int): 类别总数。
        ignore_index (int): 忽略的类别标签，比如 255 表示不关注区域。
        
    返回：
        float: 平均像素准确率 (mPA)。
    """
    # 先展平
    pred = pred.view(-1)
    target = target.view(-1)

    # 生成 mask，过滤掉 ignore_index 位置
    valid_mask = (target != ignore_index)
    pred = pred[valid_mask]
    target = target[valid_mask]

    acc = []
    for cls in range(num_classes):
        cls_pred = (pred == cls)
        cls_target = (target == cls)
        correct = (cls_pred & cls_target).sum().float()
        total = cls_target.sum().float()

        if total == 0:
            acc.append(float('nan'))
        else:
            acc.append((correct / total).item())

    return np.nanmean(acc)

def evaluate_metrics(pred, target, num_classes, ignore_index=255):
    """
    计算 mIoU 和 mPA。
    
    参数：
        pred (torch.Tensor): 预测的类别标签，形状为 (N, H, W)。
        target (torch.Tensor): 真实的类别标签，形状为 (N, H, W)。
        num_classes (int): 类别总数。
        ignore_index (int): 忽略的类别标签，比如 255 表示不关注区域。
        
    返回：
        dict: 包含 'mIoU' 和 'mPA' 的字典。
    """
    miou = calculate_iou(pred, target, num_classes, ignore_index=ignore_index)
    mpa = mean_pixel_accuracy(pred, target, num_classes, ignore_index=ignore_index)
    return {'mIoU': miou, 'mPA': mpa}


class SegmentationMetric:
    def __init__(self, num_classes, ignore_index=255):
        self.num_classes = num_classes
        self.ignore_index = ignore_index
        self.confusion_matrix = np.zeros((num_classes, num_classes), dtype=np.int64)

    def _fast_hist(self, label_true, label_pred):
        mask = (label_true != self.ignore_index) & (label_true >= 0) & (label_true < self.num_classes)
        hist = np.bincount(
            self.num_classes * label_true[mask].astype(int) + label_pred[mask],
            minlength=self.num_classes**2
        ).reshape(self.num_classes, self.num_classes)
        return hist

    def update(self, pred, target):
        if hasattr(pred, 'cpu'):
            pred = pred.cpu().numpy()
        if hasattr(target, 'cpu'):
            target = target.cpu().numpy()

        pred = pred.reshape(-1)
        target = target.reshape(-1)

        hist = self._fast_hist(target, pred)
        self.confusion_matrix += hist

        return self.evaluate()

    def evaluate(self):
        hist = self.confusion_matrix
        iou_list = []
        acc_list = []

        for i in range(self.num_classes):
            intersection = hist[i, i]
            union = hist[i, :].sum() + hist[:, i].sum() - intersection
            if union == 0:
                iou = float('nan')
            else:
                iou = intersection / union
            iou_list.append(iou)

            total = hist[i, :].sum()
            if total == 0:
                acc = float('nan')
            else:
                acc = intersection / total
            acc_list.append(acc)

        miou = np.nanmean(iou_list)
        mpa = np.nanmean(acc_list)
        return {'mIoU': miou, 'mPA': mpa}

    def clear(self):
        self.confusion_matrix = np.zeros((self.num_classes, self.num_classes), dtype=np.int64)
